import os

import cv2

from option import parse_args
from preprocessing.data_split import get_original_fake_pairs


def unaligned_videos(args, func, template: str):
    # real_fake_pairs = get_original_fake_pairs(args.root_dir, ['[0-9]', '[1-2][0-9]', '3[0-4]'])
    # real_fake_pairs = get_original_fake_pairs(args.root_dir, ['3[5-9]'])
    real_fake_pairs = get_original_fake_pairs(args.root_dir, ['4[0-9]'])
    for real_fake_pair in real_fake_pairs:
        assert len(real_fake_pair) == 3

        real_video_path = os.path.join(args.root_dir, real_fake_pair[0], f'{real_fake_pair[1]}.mp4')
        capture = cv2.VideoCapture(real_video_path)
        real_video_frames_num = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
        fake_video_path = os.path.join(args.root_dir, real_fake_pair[0], f'{real_fake_pair[2]}.mp4')
        capture = cv2.VideoCapture(fake_video_path)
        fake_video_frames_num = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
        if fake_video_frames_num != real_video_frames_num:
            func(template.format(real_fake_pair[0], real_fake_pair[1], fake_video_frames_num - real_video_frames_num))


def main():
    args = parse_args()
    assert os.path.exists(args.root_dir)
    args.save = False
    if args.save:
        with open('unaligned_videos_statistics.dat', 'w') as f:
            unaligned_videos(args, f.write, '{},{},{}\n')
    else:
        unaligned_videos(args, print, '{},{},{}')


if __name__ == '__main__':
    main()
